package com.mine.redis.util.utils.ratelimiter;

import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Collections;

/**
 * 限流器
 *
 * @author weixiao
 * @date 2024/10/31 22:21
 */
@Component
@Slf4j
public class RateLimiter {

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    private static final DefaultRedisScript<Long> RATE_LIMITER_SCRIPT;

    static {
        /*
         * ARGV[1] : 窗口的起始时间
         * ARGV[2] : 窗口内允许的最大请求数量
         * ARGV[3] : 本次请求的时间
         */
        String lua = "redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) \n" +
                "if redis.call('ZCARD', KEYS[1]) < tonumber(ARGV[2]) then \n" +
                "    redis.call('ZADD', KEYS[1], ARGV[3], ARGV[3]) \n" +
                "    return 1 \n" +
                "end \n" +
                "return 0";
        RATE_LIMITER_SCRIPT = new DefaultRedisScript<>();
        RATE_LIMITER_SCRIPT.setScriptText(lua);
        RATE_LIMITER_SCRIPT.setResultType(Long.class);
    }

    /**
     * 是否允许请求执行
     * 1.移动窗口，也就是remove掉窗口起始时间之前的key
     * 2.判断窗口内部的请求数量是否大于limitRequest
     * 2.1 如果大于，拒绝请求
     * 2.2 如果小于，将本次请求的时间戳放进zset中
     *
     * @param key redis key
     * @param windowSize 滑动窗口的大小（毫秒）
     * @param limitRequest 滑动窗口内允许的最大请求数量
     * @return 是否允许请求执行
     */
    public boolean allowRequest(String key, long windowSize, int limitRequest) {
        long currentTimeMillis = System.currentTimeMillis();
        // 窗口的起始时间
        long windowStartTimeMillis = currentTimeMillis - windowSize;
        log.info("currentTimeMillis = {}, windowStartTimeMillis = {}", currentTimeMillis, windowStartTimeMillis);
        Long result = stringRedisTemplate.execute(RATE_LIMITER_SCRIPT,
                Collections.singletonList(key),
                String.valueOf(windowStartTimeMillis), String.valueOf(limitRequest), String.valueOf(currentTimeMillis));
        return (result != null && result > 0);
    }
}
